package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractDoubleDistortionTable;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AbstractSparseTranslationTable;
import geppetto.cat.models.M1;
import geppetto.cat.models.SparseTranslationTable;
import geppetto.phraseHMM.BijectiveHMM;
import geppetto.phraseHMM.BijectiveM1;
import geppetto.phraseHMM.IBMM1;
import geppetto.phraseHMM.RegularHMM;
import geppetto.phraseHMM.SymmetricHMM;
import geppetto.phraseHMM.SymmetricM1;

import java.io.IOException;





/**
 * This class decodes the test set for a given Model receive as argument.
 * It decodes before and after projection.
 * 
 * @author javg
 *
 */
public class DecodeModels {

	/**
	 * Given a forward and a backward model of any kind, returns 
	 * 3 Pairs of models. 
	 * - Normal (HMM, M1)
	 * - Symmetric (HMM, M1)
	 * - Bijective (HMM, M1)
	 * Initialized with corresponding values. This way one can decode each pair with projection and get all
	 * combinations.
	 * @return
	 */
	public static AbstractModel[] getAllModels(AbstractModel[] originalModels){
		AbstractModel[] allmodels = new AbstractModel[6];
		if(originalModels[0].getName().contains("M1")){
			allmodels[0] = new IBMM1(originalModels[0]._corpus,(SparseTranslationTable) ((IBMM1)originalModels[0])._tb,0.0);
			allmodels[1] = new IBMM1(originalModels[1]._corpus,(SparseTranslationTable) ((IBMM1)originalModels[1])._tb,0.0);
			allmodels[2] = new BijectiveM1(originalModels[0]._corpus,(SparseTranslationTable) ((IBMM1)originalModels[0])._tb,0.0,0.001,0,1000,1000);
			allmodels[3] = new BijectiveM1(originalModels[1]._corpus,(SparseTranslationTable) ((IBMM1)originalModels[1])._tb,0.0,0.001,0,1000,1000);
			allmodels[4] = new SymmetricM1(originalModels[0]._corpus,originalModels[1]._corpus,
					(SparseTranslationTable) ((IBMM1)originalModels[0])._tb,
					(SparseTranslationTable) ((IBMM1)originalModels[1])._tb
					,0.0,0.001,0.001,1000,1000);
			allmodels[5] = ((SymmetricM1)allmodels[4]).getBackwardModel();
		}else if(originalModels[0].getName().contains("HMM")&& !originalModels[0].getName().contains("Symmetric")){
			AbstractSparseTranslationTable tbf = (AbstractSparseTranslationTable) ((RegularHMM)originalModels[0])._tb;
			AbstractSparseTranslationTable tbb = (AbstractSparseTranslationTable) ((RegularHMM)originalModels[1])._tb;
			AbstractDoubleDistortionTable tdf = (AbstractDoubleDistortionTable)((RegularHMM)originalModels[0])._distortion;
			AbstractDoubleDistortionTable tdb = (AbstractDoubleDistortionTable)((RegularHMM)originalModels[1])._distortion;
			allmodels[0] = new RegularHMM(originalModels[0]._corpus,tbf,tdf,0.0);
			allmodels[1] = new RegularHMM(originalModels[1]._corpus,tbb,tdb,0.0);
			allmodels[2] = new BijectiveHMM(originalModels[0]._corpus,tbf,tdf,0.0,0.001,0,1000,1000);
			allmodels[3] = new BijectiveHMM(originalModels[1]._corpus,tbb,tdb,0.0,0.001,0,1000,1000);
			allmodels[4] = new SymmetricHMM(originalModels[0]._corpus,originalModels[1]._corpus,
					tbf,tbb,tdf,tdb
					,0.0,0.001,0.001,1000,1000);
			allmodels[5] = ((SymmetricHMM) allmodels[4]).getBackwardModel();
		}else if(originalModels[0].getName().contains("HMM") && originalModels[0].getName().contains("Symmetric")){
			AbstractSparseTranslationTable tbf = (AbstractSparseTranslationTable) ((SymmetricHMM)originalModels[0]).forward._tb;
			AbstractSparseTranslationTable tbb = (AbstractSparseTranslationTable) ((SymmetricHMM)originalModels[1]).forward._tb;
			AbstractDoubleDistortionTable tdf = (AbstractDoubleDistortionTable)((SymmetricHMM)originalModels[0]).forward._distortion;
			AbstractDoubleDistortionTable tdb = (AbstractDoubleDistortionTable)((SymmetricHMM)originalModels[1]).forward._distortion;
			allmodels[0] = new RegularHMM(originalModels[0]._corpus,tbf,tdf,0.0);
			allmodels[1] = new RegularHMM(originalModels[1]._corpus,tbb,tdb,0.0);
			allmodels[2] = new BijectiveHMM(originalModels[0]._corpus,tbf,tdf,0.0,0.001,0,1000,1000);
			allmodels[3] = new BijectiveHMM(originalModels[1]._corpus,tbb,tdb,0.0,0.001,0,1000,1000);
			allmodels[4] = new SymmetricHMM(originalModels[0]._corpus,originalModels[1]._corpus,
					tbf,tbb,tdf,tdb
					,0.0,0.001,0.001,1000,1000);
			allmodels[5] = ((SymmetricHMM) allmodels[4]).getBackwardModel();
		}else{
			System.out.println("Unknow model leaving");
			System.exit(-1);
		}	
		return allmodels;
	}
	
	public static void getCurves(BilingualCorpus corpus, byte sentenceSource, boolean projectPosteriors, AbstractModel model, String grepString){
		int[] sentences = new int[corpus.getNumSentences(sentenceSource)];
		for (int i = 0; i < sentences.length; i++) {
			sentences[i] = i;
		}
		AlignmentsSet predicted = model.posteriorAlignments(sentences,
				sentenceSource, 1f , projectPosteriors,false);
		for(Alignment a : predicted._alignments){
			a.resetPoints();
		}
		int divisions = 100;
		for (int i = 1; i <= divisions; i++) {
			model.predictAux(predicted, 1f * i / divisions);
			float[] results = AlignmentEvaluator.calculateMeasures(predicted,
					corpus.getAlignments(sentences, sentenceSource));
			
			for(Alignment a : predicted._alignments){
				a.resetPoints();
			}
			System.out.println(grepString + " " + 1f * i / divisions + " P: " + results[0] + " R: " + results[1]);
		}
	}
	
	public static void main(String[] args) throws IOException {
		
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]);
		int maxSentenceSize = Integer.parseInt(args[2]);
		
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("Max Sent length " + maxSentenceSize);
		
		BilingualCorpus forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(corpusDescription, size,maxSentenceSize);
		BilingualCorpus backwardCorpus = forwardCorpus.reverse();	
	
		String forwardModelDir = args[3];
		String backwardModelDir = args[4];
		String modelName1 = args[5];
		

		System.out.println("Model1 ");
		System.out.println("Forward Model Dir " + forwardModelDir);
		System.out.println("Backward Model Dir " + backwardModelDir);
		System.out.println("Model Name " + modelName1);
		
		
		AbstractModel[] models= GenericMethods.initializeModelsFullDir(forwardCorpus, backwardCorpus, modelName1, 
				forwardModelDir,backwardModelDir, 
				0, 0, 0, 0);
		
		models = getAllModels(models);
		AlignmentsSet[] normal = GenericMethods.getAlignmentSet(models[0], models[1],"post-aer", true, 0);
		System.out.println("No decoding");
		System.out.println("Forward" + AlignmentEvaluator.evaluate(normal[0],forwardCorpus.getGold()));
		System.out.println("Backward" + AlignmentEvaluator.evaluate(normal[1],backwardCorpus.getGold()));
		
		getCurves(forwardCorpus, BilingualCorpus.TEST_CORPUS, false, models[0], "fnp");
		getCurves(backwardCorpus, BilingualCorpus.TEST_CORPUS, false, models[1], "bnp");
		
		
		AlignmentsSet[] bijective = GenericMethods.getAlignmentSet(models[2], models[3],"post-aer", true, 0);
		System.out.println("Bijective");
		System.out.println("Forward" + AlignmentEvaluator.evaluate(bijective[0],forwardCorpus.getGold()));
		System.out.println("Backward" + AlignmentEvaluator.evaluate(bijective[1],backwardCorpus.getGold()));
		
		getCurves(forwardCorpus, BilingualCorpus.TEST_CORPUS, true, models[2], "fbp");
		getCurves(backwardCorpus, BilingualCorpus.TEST_CORPUS, true, models[3], "bbp");
		
		AlignmentsSet[] symmetric = GenericMethods.getAlignmentSet(models[4], models[5],"post-aer", true, 0);
		System.out.println("Symmetric");
		System.out.println("Forward" + AlignmentEvaluator.evaluate(symmetric[0],forwardCorpus.getGold()));
		System.out.println("Backward" + AlignmentEvaluator.evaluate(symmetric[1],backwardCorpus.getGold()));
		
		getCurves(forwardCorpus, BilingualCorpus.TEST_CORPUS, true, models[4], "fsp");
		getCurves(backwardCorpus, BilingualCorpus.TEST_CORPUS, true, models[5], "bsp");
	}
	
	
}
